from Transformer_input import *


# 多头注意力机制类 MultiHeadedAttention 实现思路分析
# 1 init函数  (self, head, embedding_dim, dropout=0.1)
    # 每个头特征尺寸大小self.d_k  多少个头self.head  线性层列表self.linears
    # self.linears = clones(nn.Linear(embedding_dim, embedding_dim), 4)
    # 注意力权重分布self.attn=None  dropout层self.dropout
# 2 forward(self, query, key, value, mask=None)
    # 2-1 掩码增加一个维度[8,4,4] -->[1,8,4,4] 求多少批次batch_size
    # 2-2 数据经过线性层 切成8个头,view(batch_size, -1, self.head, self.d_k), transpose(1,2)数据形状变化
    #     数据形状变化[2,4,512] ---> [2,4,8,64] ---> [2,8,4,64]
    # 2-3 24个头 一起送入到attention函数中求 x, self.attn
    # attention([2,8,4,64],[2,8,4,64],[2,8,4,64],[1,8,4,4]) ==> x[2,8,4,64], self.attn[2,8,4,4]]
    # 2-4 数据形状再变化回来 x.transpose(1,2).contiguous().view(,,)
    # 数据形状变化 [2,8,4,64] ---> [2,4,8,64] ---> [2,4,512]
    # 2-5 返回最后线性层结果 return self.linears[-1](x)



# 深度copy模型 输入模型对象和copy的个数 存储到模型列表中
def clones(module,N):
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class MutiHeadAttention(nn.Module):
    def __init__(self,heads,embedding_dim,dropout=0.1):
        super().__init__()
        # 确认数据特征能否被被整除 eg 特征尺寸512 % 头数8
        assert embedding_dim % heads == 0
        # 计算每个头特征尺寸 特征尺寸512 // 头数8 = 64
        self.d_k = embedding_dim //heads
        # 多少头数
        self.heads = heads
        # 四个线性层
        self.linears = clones(nn.Linear(embedding_dim,embedding_dim),4)
        # 注意力权重分布
        self.attn = None
        # dropout层
        self.dropout = nn.Dropout(p=dropout)

    def forward(self,query,key,value,mask=None):

        # 若使用掩码，则掩码增加一个维度[8,4,4] -->[1,8,4,4]
        if mask is not None:
            mask = mask.unsqueeze(0)


        # 数据形状变化[2,4,512] ---> [2,4,8,64] ---> [2,8,4,64]
        # 4代表4个单词 8代表8个头 让句子长度4和句子特征64靠在一起 更有利捕捉句子特征
        query,key,value = [model(x).view(batch_size,-1,self.heads,self.d_k).transpose(1,2) for model,x in zip(self.linears,(query,key,value))]

        # myoutptlist_data = []
        # for model, x in zip(self.linears, (query, key, value)):
        #     print('x--->', x.shape) # [2,4,512]
        #     myoutput = model(x)
        #     print('myoutput--->',  myoutput.shape)  # [2,4,512]
        #     # [2,4,512] --> [2,4,8,64] --> [2,8,4,64]
        #     tmpmyoutput = myoutput.view(batch_size, -1,  self.head, self.d_k).transpose(1, 2)
        #     myoutptlist_data.append( tmpmyoutput )
        # mylen = len(myoutptlist_data)   # mylen:3
        # query = myoutptlist_data[0]     # [2,8,4,64]
        # key = myoutptlist_data[1]       # [2,8,4,64]
        # value = myoutptlist_data[2]     # [2,8,4,64]

        # 注意力结果表示x形状 [2,8,4,64] 注意力权重attn形状：[2,8,4,4]
        # attention([2,8,4,64],[2,8,4,64],[2,8,4,64],[1,8,4,4]) ==> x[2,8,4,64], self.attn[2,8,4,4]]
        x,self.attn = self.attention(query,key,value,mask=mask,dropout=self.dropout)

        # 数据形状变化 [2,8,4,64] ---> [2,4,8,64] ---> [2,4,512]
        x = x.transpose(1,2).reshape(batch_size,-1,self.heads*self.d_k)

        # 返回最后变化后的结果 [2,4,512]---> [2,4,512]
        x = self.linears[-1](x)

        return x


    # 自注意力机制函数attention 实现思路分析
    # attention(query, key, value, mask=None, dropout=None)
    # 1 求查询张量特征尺寸大小 d_k
    # 2 求查询张量q的权重分布socres  q@k^T /math.sqrt(d_k)
    # 形状[2,4,512] @ [2,512,4] --->[2,4,4]
    # 3 是否对权重分布scores进行 scores.masked_fill(mask == 0, -1e9)
    # 4 求查询张量q的权重分布 p_attn F.softmax()
    # 5 是否对p_attn进行dropout if dropout is not None:
    # 6 求查询张量q的注意力结果表示 [2,4,4]@[2,4,512] --->[2,4,512]
    # 7 返回q的注意力结果表示 q的权重分布
    def attention(self,query,key,value,mask=None,dropout=None):
        # query, key, value：代表注意力的三个输入张量
        # mask：代表掩码张量
        # dropout：传入的dropout实例化对象

        # 2 求查询张量q的权重分布socres  q@k^T /math.sqrt(d_k)
        # [2,4,512] @ [2,512,4] --->[2,4,4]
        scores = torch.matmul(query,key.transpose(-2,-1)/math.sqrt(self.d_k))

        # 3 是否对权重分布scores 进行 masked_fill
        if mask is not None:
            # 根据mask矩阵0的位置 对sorces矩阵对应位置进行掩码
            scores = scores.masked_fill(mask == 0 , -1e9)

        # 4 求查询张量q的权重分布 softmax
        p_attn = F.softmax(scores,dim=-1)

        # 5 是否对p_attn进行dropout
        if dropout is not None:
            p_attn = dropout(p_attn)

        # 查询张量q的注意力结果表示 bmm-matmul运算, 注意力查询张量q的权重分布p_attn
        # [2,4,4]*[2,4,512] --->[2,4,512]
        attn_putput = torch.matmul(p_attn,value)

        return attn_putput,p_attn

# 前馈全连接层 PositionwiseFeedForward 实现思路分析
# 1 init函数  (self,  d_model, d_ff, dropout=0.1):
   # 定义线性层self.w1 self.w2, self.dropout层
# 2 forward(self, x)
   # 数据经过self.w1(x) -> F.relu() ->self.dropout() ->self.w2 返回

class PositionwiseFeedForward(nn.Module):
    def __init__(self,embedding_dim,d_ff,dropout=0.1):
        # d_model  第1个线性层输入维度
        # d_ff     第2个线性层输出维度

        super().__init__()
        # 定义线性层w1 w2 dropout
        self.w1 = nn.Linear(embedding_dim,d_ff)
        self.w2 = nn.Linear(d_ff,embedding_dim)
        self.dropout = nn.Dropout(p=dropout)


    def forward(self,x):
        x = self.w1(x)
        x = F.relu(x)
        x = self.dropout(x)
        x = self.w2(x)

        return x






# 规范化层 LayerNorm 实现思路分析
# 1 init函数  (self, features, eps=1e-6):
   # 定义线性层self.a2 self.b2, nn.Parameter(torch.ones(features))
# 2 forward(self, x) 返回标准化后的结果
    # 对数据求均值 保持形状不变 x.mean(-1, keepdims=True)
    # 对数据求方差 保持形状不变 x.std(-1, keepdims=True)
    # 对数据进行标准化变换 反向传播可学习参数a2 b2
    # eg self.a2 * (x-mean)/(std + self.eps) + self.b2

class LayerNorm(nn.Module):
    def __init__(self,embedding_dim,eps = 1e-6):
        # 参数features 待规范化的数据
        # 参数 eps=1e-6 防止分母为零

        super().__init__()

        # 定义a2 规范化层的系数 y=kx+b中的k
        self.a2 = nn.Parameter(torch.ones(embedding_dim))

        # 定义b2 规范化层的系数 y=kx+b中的b
        self.b2 = nn.Parameter(torch.zeros(embedding_dim))

        self.eps = eps

    def forward(self,x):
        # 对数据求均值 保持形状不变
        # [2,4,512] -> [2,4,1]
        mean = x.mean(-1,keepdims=True)

        # 对数据求方差 保持形状不变
        # [2,4,512] -> [2,4,1]
        std = x.std(-1,keepdim=True)

        # 对数据进行标准化变换 反向传播可学习参数a2 b2
        # 注意 * 表示对应位置相乘 不是矩阵运算
        y = self.a2 * (x-mean)/(std+self.eps)+self.b2

        return y

# 子层连接结构 子层(前馈全连接层 或者 注意力机制层)+ norm层 + 残差连接
# SublayerConnection实现思路分析
# 1 init函数  (self, size, dropout=0.1):
   # 定义self.norm层 self.dropout层, 其中LayerNorm(size)
# 2 forward(self, x, sublayer) 返回+以后的结果
    # 数据self.norm() -> sublayer()->self.dropout() + x

class SublayerConnection(nn.Module):
    def __init__(self,embedding_dim,dropout=0.1):
        # 参数size 词嵌入维度尺寸大小
        # 参数dropout 置零比率

        super().__init__()
        # 定义norm层
        self.norm = LayerNorm(embedding_dim)
        # 定义dropout
        self.dropout = nn.Dropout(dropout)

    def forward(self,x,sublayer):
        # 参数x 代表数据
        # sublayer 函数入口地址 子层函数(前馈全连接层 或者 注意力机制层函数的入口地址)
        # 方式1 # 数据self.norm() -> sublayer()->self.dropout() + x
        # myres = x + self.dropout(sublayer(self.norm(x)))
        # 方式2 # 数据sublayer() -> self.norm() ->self.dropout() + x
        myres = x +self.dropout(self.norm(sublayer(x)))

        return myres